import random

# 生成两个位数为bits大质数 p q
def getTwoLargePrimes(bits):
    p = random.getrandbits(bits) | (1 << (bits-1))
    # 直到找到一个素数p
    while not test(p):
        p = random.getrandbits(bits) | (1 << (bits-1))
    q = random.getrandbits(bits) | (1 << (bits-1))
    # 直到找到一个素数q 并且 p 和 q不要太近
    while (not test(q)) or (abs(p-q) < abs((1 << bits//2))):
        q = random.getrandbits(bits) | (1 << (bits-1))  
    return p, q
    

# 快速幂取模
def fastMod(a, b, c):
    result = 1
    # 重点是将指数b做二进制分解，遍历b的每一个二进制位
    while b!= 0:
        # 当前位置是非0，进行乘模运算
        if (b&1) == 1 :
            result = (result * a) % c
        b = b >> 1
        # 每轮不断更新为a^2 mod c, a^4 mod c, ...
        a = (a * a) % c
    return result


# Miller-Rabin 素性检测算法 test(n)
def test(n):

    # 小质数加速

    primeIn100 = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

    if n in primeIn100:
        return True
    
    for prime in primeIn100:
        if n % prime == 0:
            return False

    # 找到 k 和 q，满足n = 2^k * q
    k = 0
    q = n-1
    while q%2 == 0 :
        q = q // 2
        k = k + 1

    isPrime = True # 初始化假定 n 是质数

    # 进行10轮检验
    for i in range(10):
        # 随机选择一个a，满足1<a<n-1
        a = random.randint(2, n-2)

        # 条件1判断
        # 快速幂取模算法计算 a^q mod n是否为 1 或 n-1
        if(fastMod(a, q, n)==1 or fastMod(a, q, n)==n-1):
            continue # 是质数

        # 条件1不满足，进入条件2判断
        # 快速幂取模算法计算 a^((2^j)^q) mod n 是否为 n-1
        satisfy2 = False # 初始化 条件2 不满足
        for j in range(1, k):
            if(fastMod(a, pow(2, j)*q, n)==n-1):
                satisfy2 = True # 其中一次满足即可
        
        # 满足条件2
        if satisfy2 :
            continue

        # 条件1 和 条件2均不满足
        isPrime = False
    
    return isPrime


# 欧几里得算法求两个数最大公因数
def gcd(a, b):
    # 保证 a>=b 
    if a < b: 
        a, b = b, a
    # 辗转相除法
    while a%b != 0:
        a, b = b, a%b
    return b


# 拓展的欧几里得算法 a*x + b*y = gcd(a, b)，求x，y，gcd(a, b)
def ext_gcd(a, b):
    change = False
    # 确保 a >= b
    if a < b:
        a, b = b, a
        change = True
    x1, y1, r1 = 1, 0, a
    x2, y2, r2 = 0, 1, b
    x, y, r = (x1 - (r1//r2)*x2), (y1 - (r1//r2)*y2), r1%r2
    while r != 0 :
        x1, y1, r1 = x2, y2, r2
        x2, y2, r2 = x, y, r
        x, y, r = (x1 - (r1//r2)*x2), (y1 - (r1//r2)*y2), r1%r2
    # 恢复
    if change :
        x2, y2, r2 = y2, x2, r2
    return x2, y2, r2


# 获得逆元 a * a^-1 = 1 mod b
def getInverse(a, b):
    x, y, r = ext_gcd(a, b)
    if r!=1 :
        print("获取逆元失败")
    else:
        return x % b

# 传入p和q生成公钥
def getPU(p, q):
    # 公钥n
    n = p * q
    # 欧拉函数f(n)
    fn = (p-1) * (q-1)
    # 公钥e
    i = 16 # 公钥e的指数
    e = pow(2, i) + 1
    while gcd(e, fn)!=1 : # fn 互质的e
        i = i + 1
        if(pow(2, i) + 1 < fn): # 要求e < fn
            e = pow(2, i) + 1
        else:
            print("生成公钥e失败")
            break # 出错了
    return e, n
    

# 生成私钥
def getPR(e, p, q):
    # 私钥n
    n = p * q
    # 欧拉函数f(n)
    fn = (p-1) * (q-1)
    # 私钥d
    return getInverse(e, fn), n

# 读取文件
def readFile(filename):
    fd = open(filename, 'r')
    content = fd.read()
    fd.close()
    return content

# 将字符串转为其对应的十进制ascii码的拼接形式
# 一个字符对应3位长的十进制数
def toTenAscii(string):
    tenAscii = ""
    for ch in string:
        tenAscii += str(ord(ch)).rjust(3, '0')
    return tenAscii

# 将十进制ascii码转为对应的字符串
def toString(ascii):
    tenAscii = ascii
    result = ""
    while tenAscii:
        result += chr(int(tenAscii[0:3]))
        tenAscii = tenAscii[3:]
    return result

# rsa加密
def rsa(content, e, n):
    M = content
    return fastMod(M, e, n)

# rsa解密
def deRsa(cyber, d, n):
    C = cyber
    return fastMod(C, d, n)

# 分组rsa加密
def groupEncrypt(content, c, e, n):
    cyber = ""
    tenAscii = "" 
    while content :
        if len(content) < c:
            tenAscii += toTenAscii(content[0:]).rjust(3*c, "0")
            cyber += str(rsa(int(toTenAscii(content[0:])), e, n)).rjust(len(str(n)), "0")
            content = ""
        else:
            tenAscii += toTenAscii(content[0:c])
            cyber += str(rsa(int(toTenAscii(content[0:c])), e, n)).rjust(len(str(n)), "0")
            content = content[c:]
    print("全部明文对应的10进制ascii码为: " + tenAscii)
    return cyber

def groupDecrypt(cyber, c, d, n):
    content = ""
    groupLen = len(str(n))
    while cyber :
        content += toString(str(deRsa(int(cyber[0:groupLen]), d, n)).rjust(3*c, "0"))
        cyber = cyber[groupLen:]
    return content

if __name__ == '__main__':
    print("======== RSA 公私钥生成 ========")
    # 获取两个大质数 p, q
    p, q = getTwoLargePrimes(30)
    # 获取公钥e, n
    e, n = getPU(p, q)
    print("两个大质数为: p =", p, ", q =", q)
    print("公钥为: e =", e,", n =", n)
    # 获取私钥d, n
    d, n = getPR(e, p, q)
    print("私钥为: d =", d,", n =", n)
    print("注: n包含的十进制数的长度为", len(str(n)))
    print()

    print("======== 分组 RSA 加密开始 ========")
    c = 4 # 明文每组字符数
    # 分组加密明文
    print("注: 明文分组方式为: 明文字符",c, "个一分组, 明文每四个字符的ascii码对应的", 3*c, "位十进制数作为该组RSA加密的输入数据")
    # 读取明文
    content = readFile("lab2-Plaintext.txt")
    print("读入的原明文为: ", content)
    cyber = groupEncrypt(content, c, e, n)
    print("分组RSA加密得到的密文为: " + cyber)
    print()

    print("======== 分组 RSA 解密开始 ========")
    content = groupDecrypt(cyber, c, d, n)
    print("分组RSA解密得到的明文为: " + content)
    print()

    print("======== 分组 RSA 加解密演示结束 ========\n感谢您的观看~~")
    print()
    
    

